set -x

EVAL_DATA_DIR=eval
OUTPUT_DIR=eval_output
CKPT_NAME=VideoLLaMA2-7B
CKPT=DAMO-NLP-SG/${CKPT_NAME}

gpu_list="${CUDA_VISIBLE_DEVICES:-0}"
IFS=',' read -ra GPULIST <<< "$gpu_list"

# divide data via the number of GPUs per task
GPUS_PER_TASK=1
CHUNKS=$((${#GPULIST[@]}/$GPUS_PER_TASK))

output_file=${OUTPUT_DIR}/perception_test_mcqa/answers/${CKPT_NAME}/merge.json

for IDX in $(seq 0 $((CHUNKS-1))); do
    # select the GPUs for the task
    gpu_devices=$(IFS=,; echo "${GPULIST[*]:$(($IDX*$GPUS_PER_TASK)):$GPUS_PER_TASK}")    
    TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES=${gpu_devices} python3 videollama2/eval/inference_video_mcqa_perception_test_mcqa.py \
        --model-path ${CKPT} \
        --video-folder ${EVAL_DATA_DIR}/perception_test_mcqa/videos \
        --question-file ${EVAL_DATA_DIR}/perception_test_mcqa/mc_question_test.json \
        --answer-file ${OUTPUT_DIR}/perception_test_mcqa/answers/${CKPT_NAME}/${CHUNKS}_${IDX}.json \
        --num-chunks $CHUNKS \
        --chunk-idx $IDX &
done

wait

# Clear out the output file if it exists.
> "$output_file"

echo "{" >> "$output_file"

# Loop through the indices and concatenate each file.
for IDX in $(seq 0 $((CHUNKS-1))); do
    cat ${OUTPUT_DIR}/perception_test_mcqa/answers/${CKPT_NAME}/${CHUNKS}_${IDX}.json >> "$output_file"
done

sed -i '$s/.$//' $output_file

echo "}" >> "$output_file"
